#pyuic5 -o ./qt/imgClassfy.py ./qt/imgClassfy.ui
#pyuic5 -o ./qt/video2img.py ./qt/video2img.ui
import sys
from PyQt5.QtWidgets import QApplication, QMainWindow,QFileDialog,QDesktopWidget
from PyQt5.QtWidgets import QMessageBox
from PyQt5.QtGui import QImage,QPixmap,QTransform
from qt.imgClassfy import Ui_MainWindow
from qt.save_paths import Ui_saving_path
from qt.video2img import Ui_Video2imgWindow
from shutil import copyfile

import os
# import threading
import logging
import time
logging.basicConfig(filename='./Classfy_Img.log', filemode='a', level=logging.DEBUG,
                    format='%(asctime)s-%(levelname)s : %(message)s')
logging.info('Classfy_Img_v1.1_20200912')
## 界面主框架
class MainWindow(QMainWindow,Ui_MainWindow):
    def __init__(self):
        super(MainWindow,self).__init__()
        # 初始化参数
        self.imgs_path = []
        self.img_name = []
        self.origin_img = QPixmap()
        self.resize_img = QPixmap()
        self.img_num = 0
        self.img_resize_ratio = 0
        self.img_class = {}
        # 类别数量统计
        self.class_num = [0,0,0,0,0]
        # 保存路径
        self.save_path = ['./classfy/0','./classfy/1','./classfy/2','./classfy/3','./classfy/4']
        self.makedir()
        # 初始化UI
        self.setupUi(self)
        self.center()
        self.frameRate = 20
        # 按键使能
        self.ac_LastImg.setEnabled(False)
        self.ac_nextImg.setEnabled(False)
        self.ac_class_0.setEnabled(False)
        self.ac_class_1.setEnabled(False)
        self.ac_class_2.setEnabled(False)
        self.ac_class_3.setEnabled(False)
        self.ac_class_4.setEnabled(False)
        # self.stopEvent = threading.Event()
        # self.stopEvent.clear()
        ###  控件 连接 信号与槽
        #         ****命令行控件****
        self.imgFile.triggered.connect(self.get_imgFiles)
        self.ac_LastImg.triggered.connect(self._get_last_img)
        self.ac_nextImg.triggered.connect(self._get_next_img)
        self.ac_class_0.triggered.connect(self._set_class_0)
        self.ac_class_1.triggered.connect(self._set_class_1)
        self.ac_class_2.triggered.connect(self._set_class_2)
        self.ac_class_3.triggered.connect(self._set_class_3)
        self.ac_class_4.triggered.connect(self._set_class_4)
        self.ac_save_paths.triggered.connect(self._saving_paths)
        self.ac_video2img.triggered.connect(self._video2img)
        # 保存文件路径 窗口
        self.save_path_wimdow = SavingPath()
        self.save_path_wimdow.btn_path_0.clicked.connect(self._saving_path_0)
        self.save_path_wimdow.btn_path_1.clicked.connect(self._saving_path_1)
        self.save_path_wimdow.btn_path_2.clicked.connect(self._saving_path_2)
        self.save_path_wimdow.btn_path_3.clicked.connect(self._saving_path_3)
        self.save_path_wimdow.btn_path_4.clicked.connect(self._saving_path_4)
        #视频抽帧窗口
        self.video2img_window = Vedio2Image()
    def makedir(self):
        if not os.path.exists('./classfy'):
            os.makedirs('./classfy')
        for path in self.save_path:
            if not os.path.exists(path): os.makedirs(path)
    # 窗口居中
    def center(self):

        # QtCore.QCoreApplication.setAttribute(QtCore.Qt.AA_EnableHighDpiScaling)
        #获得窗口
        framePos = self.frameGeometry()
        #获得屏幕中心点
        scPos  = QDesktopWidget().availableGeometry().center()
        # screen = QDesktopWidget().screenGeometry()
        # self.resize(screen.width(), screen.height())
        #设置固定尺寸
        # self.setFixedSize(screen.width(), screen.height()-100)
        # 显示到屏幕中心
        framePos.moveCenter(scPos)
        self.move(framePos.topLeft())
        ### 获取控件尺寸与位置
        # print(self.lb_img_show.width())
        # print(self.lb_img_show.geometry())
        # print(self.lb_img_show.frameGeometry())
    # 窗口关闭事件
    def closeEvent(self, event):
        # self.stopEvent.set()
        logging.info("classify num : {}".format(str(self.img_num + 1) + '/'+ str(len(self.imgs_path))))
        logging.info("=================Exit================== ")
        # 关闭子窗口
        reply = QMessageBox.question(self, 'EXIT',
            "Are you sure to quit?", QMessageBox.Yes | 
            QMessageBox.No, QMessageBox.Yes)
        if reply == QMessageBox.Yes:           
            event.accept()       
        else:
            event.ignore()   
    # 【1. 获取文件绝对路径】
    def _getFilesPath(self,file_dir):
        filenames = []
        img_name = []
        for root, dirs, files in os.walk(file_dir,topdown=False):
            for name in files:
                # 只选择bmp jpg 
                # print(name.split('.')[-1])
                if name.split('.')[-1] in ['bmp','jpg','jpeg','png']:
                    # print(os.path.join(root, name))
                    filenames.append(os.path.join(root, name))
                    img_name.append(name)
        return filenames,img_name

    def _save_img_init(self):
        # 图片分类初始化 查询目标文件夹中是否已有分类，有则读入字典中
        pass
    def _get_last_img(self):
        self.img_num  = self.img_num - 1 if self.img_num > 0 else len(self.imgs_path)-1
        self._img_display()
    def _get_next_img(self):
        self.img_num = self.img_num + 1 if self.img_num < len(self.imgs_path)-1 else 0
        self._img_display()
        if self.img_num == len(self.imgs_path)-1:
            QMessageBox.warning(self, 'Tips','Last one!') 
    def _img_classfy(self,class_n):
        self.class_num[class_n] = self.class_num[class_n] + 1
        # print(self.img_name[self.img_num])
        if self.imgs_path[self.img_num] in self.img_class.keys():
            # 如果该图片已经分类 则需要重新分类 将之前分类图删除 并重新保存到新文件夹
            last_class = self.img_class[self.imgs_path[self.img_num]]
            self.class_num[last_class] = self.class_num[last_class] - 1 
            os.remove(self.save_path[last_class] + "/" + self.img_name[self.img_num])
        self.img_class[self.imgs_path[self.img_num]] = class_n
        # print(self.imgs_path[self.img_num],self.save_path[class_n])
        copyfile(self.imgs_path[self.img_num],self.save_path[class_n]+ "/" + self.img_name[self.img_num])
        self._get_next_img()
    def _set_class_0(self): self._img_classfy(0)
    def _set_class_1(self): self._img_classfy(1)
    def _set_class_2(self): self._img_classfy(2)
    def _set_class_3(self): self._img_classfy(3)
    def _set_class_4(self): self._img_classfy(4)


    def get_imgFiles(self):
        # self.stopEvent.set()
        # logging.info('runing thread : {}'.format(threading.enumerate()))   
        """
        单个文件打开 QFileDialog.getOpenFileName()
        多个文件打开 QFileDialog.getOpenFileNames()
        文件夹选取 QFileDialog.getExistingDirectory()
        文件保存 QFileDialog.getSaveFileName()
        """
        img_files = QFileDialog.getExistingDirectory(None,"Images dir","./") #起始路径
        if img_files.strip():
            logging.debug('open the images file...')
            self.video_name = ''
            logging.debug('images file is {}'.format(img_files))
            self.imgs_path,self.img_name = self._getFilesPath(img_files)
            self._img_display()
            self.ac_LastImg.setEnabled(True)
            self.ac_nextImg.setEnabled(True)
            self.ac_class_0.setEnabled(True)
            self.ac_class_1.setEnabled(True)
            self.ac_class_2.setEnabled(True)
            self.ac_class_3.setEnabled(True)
            self.ac_class_4.setEnabled(True)
        else:
            logging.warning("Open image file error!!!")
            QMessageBox.warning(self, 'Warning','images file is empty!')



    def _get_min(self,x1,x2):
        return x1 if x1 < x2 else x2
    def _img_display(self):
        if len(self.imgs_path)>0:
            self.origin_img.load(self.imgs_path[self.img_num])
            # self.lb_img_show.setPixmap(self.origin_img)
            self.img_resize_ratio = self._get_min(self.lb_img_show.width()/self.origin_img.width(),self.lb_img_show.height()/self.origin_img.height())
            self.resize_img = self.origin_img.scaled(int(self.origin_img.width()*self.img_resize_ratio),int(self.origin_img.height()*self.img_resize_ratio ))
            self.lb_img_show.setPixmap(self.resize_img)
        # 显示图像
        # self._label_show_img(self.origin_img)
        # self.origin_img.load(self.imgs_path[self.img_num])
        # 显示分类信息 图片数量 分类结果
        self.ed_img_num.setText(str(self.img_num + 1) + '/'+ str(len(self.imgs_path)))
        temp_class = str(self.img_class[self.imgs_path[self.img_num]]) if self.imgs_path[self.img_num] in self.img_class.keys() else None
        self.ed_img_class.setText(temp_class)
        self.ed_class_0.setText(str(self.class_num[0]))
        self.ed_class_1.setText(str(self.class_num[1]))
        self.ed_class_2.setText(str(self.class_num[2]))
        self.ed_class_3.setText(str(self.class_num[3]))
        self.ed_class_4.setText(str(self.class_num[4]))
    ### ==================路径保存窗口===========================
    def _saving_paths(self):
        self.save_path_wimdow.show()
        pass
    def _saving_path_0(self):
        tmp_save_file = QFileDialog.getExistingDirectory(None,"Save dir","./") 
        self.save_path_wimdow.ed_path_0.setText(tmp_save_file)
        self.save_path[0] = tmp_save_file
        logging.info("Class 0 save path = {}".format(tmp_save_file))
    def _saving_path_1(self):
        tmp_save_file = QFileDialog.getExistingDirectory(None,"Save dir","./") 
        self.save_path_wimdow.ed_path_1.setText(tmp_save_file)
        self.save_path[1] = tmp_save_file
        logging.info("Class 1 save path = {}".format(tmp_save_file))
    def _saving_path_2(self):
        tmp_save_file = QFileDialog.getExistingDirectory(None,"Save dir","./") 
        self.save_path_wimdow.ed_path_2.setText(tmp_save_file)
        self.save_path[2] = tmp_save_file
        logging.info("Class 2 save path = {}".format(tmp_save_file))
    def _saving_path_3(self):
        tmp_save_file = QFileDialog.getExistingDirectory(None,"Save dir","./") 
        self.save_path_wimdow.ed_path_3.setText(tmp_save_file)
        self.save_path[3] = tmp_save_file
        logging.info("Class 3 save path = {}".format(tmp_save_file))
    def _saving_path_4(self):
        tmp_save_file = QFileDialog.getExistingDirectory(None,"Save dir","./") 
        self.save_path_wimdow.ed_path_4.setText(tmp_save_file)
        self.save_path[4] = tmp_save_file
        logging.info("Class 4 save path = {}".format(tmp_save_file))
    def _video2img(self):
        self.video2img_window.show()


## 保存路径
class SavingPath(QMainWindow,Ui_saving_path):
    def __init__(self):
        super(SavingPath,self).__init__()
        self.setupUi(self)
# 视频抽帧
import cv2
import threading
class Vedio2Image(QMainWindow,Ui_Video2imgWindow):
    
    def __init__(self):
        super(Vedio2Image,self).__init__()
        self.setupUi(self)
        self.video_name = ''
        self.img_files = "./"
        self._img_exname = ''
        self._vedio2img_count = 10
        self.ed_img_count.setText('50')
        self.ed_img_exname.setText('test')
        # 
        self.btn_open_video.clicked.connect(self._open_video)
        self.btn_save_path.clicked.connect(self._img_save_path)
        self.btn_start.clicked.connect(self._start_video2img)
        self.btn_stop.clicked.connect(self._stop_video2img)
        #
        self.stopEvent = threading.Event()
        self.stopEvent.clear()

    def _open_video(self):
        video_path,fileType = QFileDialog.getOpenFileName(None,"Video file","./") #起始路径   
        # print(video_path[0])
        self.video_name = video_path if video_path.split('.')[-1] in ['mp4','avi','h264'] else self.video_name
        self.ed_video_name.setText(self.video_name)
    def _img_save_path(self):
        self.img_files = QFileDialog.getExistingDirectory(None,"Images dir",self.img_files) #起始路径
        self.ed_save_path.setText(self.img_files)
    def VideoExtract(self):
        videoCapture = cv2.VideoCapture(self.video_name)
        # fps = videoCapture.get(cv2.CAP_PROP_FPS)
        # 获取视频总帧数
        frame_all = videoCapture.get(cv2.CAP_PROP_FRAME_COUNT)
        img_num = int(frame_all/self._vedio2img_count)
        # print("total fps = {}".format(frame_all))
        i = 0 
        j = 0
        while True:
            success, frame = videoCapture.read()
            i += 1
            if (i % self._vedio2img_count == 0):
                # 保存图片
                j += 1
                savedname = self.img_files + '\\' + self._img_exname + str(j) + '.bmp'
                cv2.imwrite(savedname, frame)
                self.pb_img_num.setValue(int(j/img_num * 100))
                # print('image of %s is saved' % (savedname))
            if not success:
                logging.error('video is not open!')
                break
            if True == self.stopEvent.is_set():
                logging.info("video to img : {}/{}".format(i,frame_all))
                logging.info('=================stop video2img thread============')
                break
            
        # QMessageBox.warning(self, 'Tips','video to image is over! ') 
    def _start_video2img(self):
        self._img_exname = self.ed_img_exname.text()
        self._vedio2img_count = int(self.ed_img_count.text())
        logging.info('exname : {} \n img_count : {}'.format(self._img_exname,self._vedio2img_count))
        self.stopEvent.clear()
        th_img = threading.Thread(target=self.VideoExtract,name='Video2ImageThread')
        th_img.start()
        logging.info('start video2img threading = {}'.format(threading.enumerate()))
    def _stop_video2img(self):
        self.stopEvent.set()
        self.pb_img_num.setValue(0)
        pass
if __name__ == "__main__":
    app = QApplication(sys.argv)
    Gui = MainWindow()
    Gui.show()
    sys.exit(app.exec_())